use std::{ffi::{CString, OsStr}, path::Path};
use libc::{chdir, chroot, execv};
use nix::{libc::sethostname, mount::{mount, umount, MsFlags}, sched::{unshare, CloneFlags}, sys::wait::waitpid, unistd::{fork, getppid, ForkResult}};
use tracing::{error, info};
use super::{namespace::setup_user_namespace, tools::ensure_directory_exists, overlayfs::setup_overlayfs};

const MERGED: &str = "rootfs/merged";

pub fn container_run() -> anyhow::Result<()> {
    match unsafe { fork() } {
        Ok(ForkResult::Parent { child }) => {
            waitpid(child, None)?;
            info!("umount {}...", MERGED);
            umount(MERGED)?;
        },
        Ok(ForkResult::Child) => {
            info!("Parent (PID: {}) created a child (PID: {})..", getppid(), std::process::id());
            {
                let lowerdir = "rootfs/lowdir";
                let upperdir = "rootfs/upperdir";
                let workdir = "rootfs/workdir";
                let procdir = "rootfs/lowdir/proc";

                ensure_directory_exists(lowerdir)?;
                ensure_directory_exists(upperdir)?;
                ensure_directory_exists(workdir)?;
                ensure_directory_exists(procdir)?;
                setup_overlayfs(lowerdir, upperdir, workdir, MERGED)?;
            }

            unshare(CloneFlags::CLONE_NEWUTS | CloneFlags::CLONE_NEWUSER | CloneFlags::CLONE_NEWPID | CloneFlags::CLONE_NEWNS)
            .expect("Failed to unshare namespaces");

            info!("Enter container {}", std::process::id());

            {
                setup_user_namespace(std::process::id())?;
                set_container_hostname("mtdocker");
            }

            {
                match unsafe{ fork() } {
                    Ok(ForkResult::Parent { child }) => {
                        info!("container process (PID: {}) exit..", child);
                        waitpid(child, None)?;
                    },
                    Ok(ForkResult::Child) => {
                        info!("enter container..");
                        setup_container()?;
                        run_container_process();
                    },
                    Err(_) => panic!("Fork failed!"),
                }
            }
        },
        Err(_) => panic!("Fork failed!"),
    }
    Ok(())
}

fn run_container_process() {
    let bash = CString::new("/bin/bash").expect("CString::new failed");
    let args = [bash.as_ptr(), std::ptr::null()];

    let res = unsafe {
        execv(bash.as_ptr(), args.as_ptr())
    };

    if res != 0 {
        panic!("execv failed!");
    }
}

fn set_container_hostname(hostname: &str) {
    let cname = CString::new(hostname).expect("CString::new failed");
    let result = unsafe { 
        sethostname(cname.as_ptr(), cname.as_bytes().len()) 
    };
    if result != 0 {
        error!("set container hostname {} error!", hostname);
    } else {
        info!("set container hostname {} success!", hostname);
    }
}

fn setup_container() -> anyhow::Result<()> {
    info!("Setting up container environment...");

    let mut dir = CString::new(MERGED).expect("CString::new failed");
    if unsafe { chdir(dir.as_ptr()) } != 0 {
        panic!("chdir {} failed!", MERGED);
    }

    dir = CString::new(".").expect("CString::new failed");
    if unsafe { chroot(dir.as_ptr()) } != 0 {
        panic!("chroot {} failed!", MERGED);
    }

    let res = mount(
        Some(OsStr::new("proc")), 
        Path::new("/proc"), 
        Some(OsStr::new("proc")),
        MsFlags::empty(),
        Some(OsStr::new("")),
    );

    match res {
        Ok(_) => info!("Proc setup completed successfully."),
        Err(e) => return Err(anyhow::anyhow!("Failed to setup proc: {}", e)),
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_cmd_execute() {
        let output = std::process::Command::new("/bin/bash").output().expect("cmd exec error!");
        println!("{}", String::from_utf8_lossy(&output.stdout));
    }
}